import numpy as np
import faiss
import struct


def ivecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('i' * len(x), *x))

    f.close()


def ip_gnd(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatIP(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    return gnd_idx, gnd_distance


if __name__ == '__main__':
    base = np.loadtxt('/home/zhengbian/mips-graph/mobius/base.txt').astype(np.float32)
    query = np.loadtxt('/home/zhengbian/mips-graph/mobius/test.txt').astype(np.float32)
    gnd_idx, gnd_dist = ip_gnd(base, query, 10)
    ivecs_write('/home/zhengbian/mips-graph/mobius/gnd.ivecs', gnd_idx)
